from flask import Flask, request, jsonify, render_template
import torch
from PIL import Image
import io
import base64
import numpy as np

app = Flask(__name__)

# Load YOLOv5 model
model = torch.hub.load('..', 'custom', path='../yolov5s.pt', source='local')

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_image():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400

    try:
        # Read image and perform inference
        image_bytes = file.read()
        img = Image.open(io.BytesIO(image_bytes)).convert('RGB')  # Convert image to RGB
        results = model(img)

        # Convert results to JSON
        results_json = results.pandas().xyxy[0].to_json(orient='records')

        # Prepare the base64-encoded result image
        result_img_array = results.render()[0]  # Render the results and get numpy array
        result_img = Image.fromarray(result_img_array)  # Convert numpy array to PIL image
        buffered = io.BytesIO()
        result_img.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()

        return jsonify({'image': img_str, 'detections': results_json})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True)
